#!/bin/bash

# 设置TensorRT路径
export TENSORRT_ROOT=/data2/xd/TensorRT-8.5.1.7.Linux.x86_64-gnu.cuda-11.8.cudnn8.6/TensorRT-8.5.1.7
export PATH="$PATH:$TENSORRT_ROOT/bin"
export LD_LIBRARY_PATH="$TENSORRT_ROOT/lib:$LD_LIBRARY_PATH"

# 设置插件路径
PLUGIN_PATH="/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/grid_sampler/build/libtrtgrid_sampler.so"

echo "=== 构建GridSample TensorRT引擎 ==="
echo "ONNX文件: gridsample_test_model.onnx"
echo "插件路径: $PLUGIN_PATH"

# 构建FP16引擎
echo "构建FP16引擎..."
trtexec \
    --onnx=gridsample_test_model.onnx \
    --saveEngine=gridsample_fp16.plan \
    --memPoolSize=workspace:2000 \
    --fp16 \
    --plugins="$PLUGIN_PATH" \
    --noBuilderCache

if [ $? -eq 0 ]; then
    echo "✅ FP16引擎构建成功!"
    echo "引擎文件: gridsample_fp16.plan"
else
    echo "❌ FP16引擎构建失败!"
fi

# 构建FP32引擎
echo "构建FP32引擎..."
trtexec \
    --onnx=gridsample_test_model.onnx \
    --saveEngine=gridsample_fp32.plan \
    --memPoolSize=workspace:2000 \
    --plugins="$PLUGIN_PATH" \
    --noBuilderCache

if [ $? -eq 0 ]; then
    echo "✅ FP32引擎构建成功!"
    echo "引擎文件: gridsample_fp32.plan"
else
    echo "❌ FP32引擎构建失败!"
fi

echo "=== 引擎构建完成 ==="
echo "检查生成的引擎文件..."
ls -la gridsample_*.plan 2>/dev/null || echo "没有生成引擎文件" 